package com.maiko.maikoaiagent.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * 查询重写器组件
 *
 * 该类利用 ChatModel 和 RewriteQueryTransformer 实现查询的自动重写功能。
 *
 * 主要用途：
 * 对输入的查询语句进行改写，使查询更精准或更符合下游检索/生成模型的需求，
 * 提升搜索或对话系统的表现效果。
 *
 * 设计说明：
 * - 通过构造函数注入 ChatModel 实例
 * - 创建 ChatClient.Builder 以驱动 ChatModel
 * - 使用 RewriteQueryTransformer 进行查询重写
 *
 * 调用流程：
 * 传入原始查询字符串，返回改写后的查询文本。
 *
 * @author Maiko7
 * @create: 2025-07-09 10:29
 */
@Component
public class QueryRewriter {

    /**
     * 查询转换器，负责执行查询重写
     */
    private final QueryTransformer queryTransformer;

    /**
     * 构造方法，注入 ChatModel 并创建 RewriteQueryTransformer 实例
     *
     * @param dashscopeChatModel 用于驱动查询重写的聊天模型
     */
    public QueryRewriter(ChatModel dashscopeChatModel) {
        // 使用注入的 ChatModel 创建 ChatClient.Builder
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);

        // 创建基于 ChatClient 的查询重写转换器
        queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
    }

    /**
     * 执行查询重写操作
     *
     * @param prompt 原始查询字符串
     * @return 重写后的查询字符串
     */
    public String doQueryRewrite(String prompt) {
        // 构造 Query 对象
        Query query = new Query(prompt);

        // 通过转换器执行查询重写
        Query transformedQuery = queryTransformer.transform(query);

        // 返回重写后的查询文本
        return transformedQuery.text();
    }
}


